// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/compression/paged_output_steam.h>

namespace kumo::pollux::dwio::common::compression {
    std::vector<melon::StringPiece> PagedOutputStream::createPage() {
        auto origSize = buffer_.size();
        POLLUX_CHECK_GT(origSize, pageHeaderSize_);
        origSize -= pageHeaderSize_;

        auto compressedSize = origSize;
        // Applies compression if there is compressor and original data size exceeds
        // threshold.
        if (compressor_ && origSize >= threshold_) {
            compressionBuffer_ = pool_->getBuffer(buffer_.size());
            compressedSize = compressor_->compress(
                buffer_.data() + pageHeaderSize_,
                compressionBuffer_->data() + pageHeaderSize_,
                origSize);
        }

        melon::StringPiece compressed;
        if (compressedSize >= origSize) {
            // write orig
            writeHeader(buffer_.data(), origSize, true);
            compressed = melon::StringPiece(buffer_.data(), origSize + pageHeaderSize_);
        } else {
            // write compressed
            writeHeader(compressionBuffer_->data(), compressedSize, false);
            compressed = melon::StringPiece(
                compressionBuffer_->data(), compressedSize + pageHeaderSize_);
        }

        if (encryptor_ == nullptr) {
            return {compressed};
        }

        encryptionBuffer_ = encryptor_->encrypt(melon::StringPiece(
            compressed.begin() + pageHeaderSize_, compressed.end()));
        updateSize(
            const_cast<char *>(compressed.begin()), encryptionBuffer_->length());
        return {
            melon::StringPiece(compressed.begin(), pageHeaderSize_),
            melon::StringPiece(
                reinterpret_cast<const char *>(encryptionBuffer_->data()),
                encryptionBuffer_->length())
        };
    }

    void PagedOutputStream::writeHeader(
        char *buffer,
        size_t compressedSize,
        bool original) {
        POLLUX_CHECK_LT(compressedSize, 1 << 23);
        buffer[0] = static_cast<char>((compressedSize << 1) + (original ? 1 : 0));
        buffer[1] = static_cast<char>(compressedSize >> 7);
        buffer[2] = static_cast<char>(compressedSize >> 15);
    }

    void PagedOutputStream::updateSize(char *buffer, size_t compressedSize) {
        POLLUX_CHECK_LT(compressedSize, 1 << 23);
        buffer[0] = ((buffer[0] & 0x01) | static_cast<char>(compressedSize << 1));
        buffer[1] = static_cast<char>(compressedSize >> 7);
        buffer[2] = static_cast<char>(compressedSize >> 15);
    }

    void PagedOutputStream::resetBuffers() {
        // Reset compression buffer size and return.
        if (compressionBuffer_ != nullptr) {
            pool_->returnBuffer(std::move(compressionBuffer_));
        }
        encryptionBuffer_ = nullptr;
    }

    uint64_t PagedOutputStream::flush() {
        const auto size = buffer_.size();
        const auto originalSize = bufferHolder_.size();
        if (size > pageHeaderSize_) {
            auto buffers = createPage();
            const auto cleanup = melon::makeGuard([this]() {
                resetBuffers();
                // Reset input buffers. clear() forces the buffer to shrink.
                // Not doing so lead to very high flush memory overhead.
                buffer_.clear();
                buffer_.resize(pageHeaderSize_);
            });
            bufferHolder_.take(std::move(buffers));
        }
        return bufferHolder_.size() - originalSize;
    }

    void PagedOutputStream::BackUp(int32_t count) {
        if (count > 0) {
            POLLUX_CHECK_GE(buffer_.size(), count + pageHeaderSize_);
            BufferedOutputStream::BackUp(count);
        }
    }

    bool PagedOutputStream::Next(void **data, int32_t *size, uint64_t increment) {
        if (!tryResize(data, size, pageHeaderSize_, increment)) {
            auto buffers = createPage();
            const auto cleanup = melon::makeGuard([this]() { resetBuffers(); });
            flushAndReset(data, size, pageHeaderSize_, std::move(buffers));
        }
        return true;
    }

    void PagedOutputStream::recordPosition(
        PositionRecorder &recorder,
        int32_t bufferLength,
        int32_t bufferOffset,
        int32_t strideIndex) const {
        // add compressed size, then uncompressed
        recorder.add(bufferHolder_.size(), strideIndex);
        auto size = buffer_.size();
        if (size) {
            size -= (pageHeaderSize_ + bufferLength - bufferOffset);
        }
        recorder.add(size, strideIndex);
    }
} // namespace kumo::pollux::dwio::common::compression
